#include "axle/core.h"

#if defined(SYS_IS_WINDOWS) && defined(BUILD_SHARED)
#define LIBAPI __declspec(dllexport)
#endif

#include "image_db.h"

#include <stdlib.h>
#include <time.h>

#include "axle/core/io.h"

#include "haar.h"

namespace oe {
using namespace std;

IntGuid::IntGuid() : next_id_(0) { }

int IntGuid::NewId() {   
  int id = next_id_++;
  // save next_id_ to config file
  return id;  
}

////////////////////////////////////////////////////////////////////
ImageDbServer::ImageDbServer() { ImageDb::Initialize(); }

int ImageDbServer::NewDb(const std::string &name) {
  int id = s_guid.NewId();
  if (id >= 0) {
    dbs_[id] = ImageDbPtr(new ImageDb(id, name), DbDeleter);
  }
  return id;
}

ImageDbServer::~ImageDbServer() { }

ImageDbServerPtr ImageDbServer::s_server = ImageDbServerPtr(new ImageDbServer, ImageDbServer::ServerDeleter);
IntGuid ImageDbServer::s_guid;

///////////////////////////////////////////////////////////////////////////

const float ImageDb::s_weights[2][6][3] = {
  // For scanned picture (sketch=0):
  //    Y      I      Q       idx total occurs
  {{ 5.00f, 19.21f, 34.37f},  // 0   58.58      1 (`DC' component)
	  { 0.83f,  1.26f,  0.36f},  // 1    2.45      3
	  { 1.01f,  0.44f,  0.45f},  // 2    1.90      5
	  { 0.52f,  0.53f,  0.14f},  // 3    1.19      7
	  { 0.47f,  0.28f,  0.18f},  // 4    0.93      9
	  { 0.30f,  0.14f,  0.27f}}, // 5    0.71      16384-25=16359

	// For handdrawn/painted sketch (sketch=1):
	//    Y      I      Q
	{{ 4.04f, 15.14f, 22.62f},
		{ 0.78f,  0.92f,  0.40f},
		{ 0.46f,  0.53f,  0.63f},
		{ 0.42f,  0.26f,  0.25f},
		{ 0.41f,  0.14f,  0.15f},
		{ 0.32f,  0.07f,  0.38f}}
};

byte ImageDb::s_img_bin[kImgSize];
bool ImageDb::s_initialized = false;

typedef ImageDb::QueryParam QueryParam;

void ImageDb::Initialize() {
  s_initialized = true;
	srand((unsigned)time(0));

	/* setup initial fixed weights that each coefficient represents */
	int i, j;

	/*
	0 1 2 3 4 5 6 i
	0 0 1 2 3 4 5 5
	1 1 1 2 3 4 5 5
	2 2 2 2 3 4 5 5
	3 3 3 3 3 4 5 5
	4 4 4 4 4 4 5 5
	5 5 5 5 5 5 5 5
	5 5 5 5 5 5 5 5
	j
	 */

	/* Every position has value 5, */
  memset(s_img_bin, 5, kImgSize);

	/* Except for the 5 by 5 upper-left quadrant: */
	for (i = 0; i < 5; i++)
		for (j = 0; j < 5; j++)
			s_img_bin[i * kDim + j] = max(i, j);
	// Note: imgBin[0] == 0
}

IntGuid ImageDb::s_guid;

ImageDb::ImageDb(int id, const std::string &name)
    : id_(id), name_(name), haar_(kDim, kNCoefs) {
  /*for (int i = 0; i < 3; ++i)
    A_[i] = new float[kImgSize];  */
}

ImageDb::~ImageDb() {
  /*for (int i = 0; i < 3; ++i)
    delete[] A_[i];*/
  /*for (SigMapIt it = sigs_.begin(); it != sigs_.end(); ++it) {
    delete it->second;
  }*/
}

IntId ImageDb::AddImage(const std::string &filename) {
  ax::Logger::Debug("AddImage: %s", filename.c_str());
  ax::ImageRgbByte img;
  if (ax::kOK != img.Load(filename)) {
    ax::Logger::Debug("Loading image failed");
    return kBadId;
  }
  
  SigPtr sig = this->ExtractSignature(img);
  sigs_[sig->id] = sig;

  // bloom filter

  // !!! 0, 1, 2 are magic numbers, should be handled  
 
  for (int i = 0; i < kNCoefs; ++i) {
    for (int c = 0; c < 3; ++c) {
      if (sig->sigs[c][i] > 0) 
        img_buckets_[c][0][sig->sigs[c][i]].push_back(sig->id);
      if (sig->sigs[c][i] < 0) 
        img_buckets_[c][1][-sig->sigs[c][i]].push_back(sig->id);
    }
  }

  images_[sig->id] = filename;

  ax::Logger::Debug("AddImage, succeed.");
  return sig->id;
}

int ImageDb::AddImages(const std::vector<std::string> &filenames,
                       std::vector<IntId> *ids) {
  ax::Logger::Debug("AddImages, %d images to be added...", filenames.size());
  for (size_t i = 0; i < filenames.size(); ++i) {
    IntId img_id = this->AddImage(filenames[i]);
    if (NULL != ids) ids->push_back(img_id);
  }
  ax::Logger::Debug("AddImages, %d imageds added.", images_.size());
  return 0;
}

int ImageDb::AddImageFolder(const std::string &dir) {
  std::vector<std::string> filenames;
  ax::GetFileNames(dir, filenames);
  return this->AddImages(filenames, NULL);  
}

int ImageDb::Query(const std::string &filename, const QueryParam &param,
                   ResultList &result) {
  ax::ImageRgbByte image;
  V_RET(image.Load(filename));

  return Query(image, param, result);
}

int ImageDb::Query(IntId img_id, const QueryParam &param, ResultList &result) {
  const SigPtr sig = this->GetSignature(img_id);
  if (NULL == sig) return 0;

  return this->Query(*sig, param, result);
}

int ImageDb::Query(const ax::ImageRgbByte &image, const QueryParam &param,
                   ResultList &result) {  
    
  SigPtr sig = this->ExtractSignature(image);
  if (NULL == sig) return 0;

  return this->Query(*sig, param, result);
}

int ImageDb::Query(const ImageDb::Signature &sig, const QueryParam &param, 
                   ResultList &result) {
                     // compute score for each image

  int sketch = param.sketch;
  for (SigMapIt it = sigs_.begin(); it != sigs_.end(); ++it) {
    it->second->score = 0;
    for (int c = 0; c < 3; ++c) {
      it->second->score += s_weights[sketch][0][c] * 
                           fabs(it->second->avgs[c] - sig.avgs[c]);
    }
  }

  if (!param.color_only) {
    for (int i = 0; i < kNCoefs; ++i) {
      for (int c = 0; c < 3; ++c) {
        int pn = 0;
        int idx = 0;
        if (sig.sigs[c][i] > 0) {
          pn = 0;
          idx = sig.sigs[c][i];
        } else {
          pn = 1;
          idx = -sig.sigs[c][i];
        }

        ImgIdListIt it = img_buckets_[c][pn][idx].begin();
        for ( ; it != img_buckets_[c][pn][idx].end(); ++it) {
          sigs_[*it]->score -= s_weights[sketch][s_img_bin[idx]][c];
        }
      }
    }
  }
  
  SigMapIt it = sigs_.begin();
  ResultPriorityQueue q;
  for (int i = 0; it != sigs_.end() && i < param.n_results; ++it, ++i) {
    q.push(Result(it->first, it->second->score));
  }

  for (; it != sigs_.end(); ++it) {
    if ((it->second->score < kTerribleScore) &&
        (it->second->score < q.top().score)) {
      q.pop();
      q.push(Result(it->first, it->second->score));
    }
  }

  while (!q.empty()) {
    result.push_back(q.top());
    q.pop();
  }

  return result.size();
}

ImageDb::SigPtr ImageDb::ExtractSignature(const ax::ImageRgbByte &image) {
  // prepare data for haar transform
  ax::ImageRgbByte resized_img;
  resized_img.Copy(image);
  resized_img.Resize(kDim, kDim);

  typedef ax::ImageRgbByte::PixelType Pixel;
  
  int idx = 0;

  for (int i = 0; i < kDim; ++i) {
    for (int j = 0; j < kDim; ++j) {
      const Pixel &p = resized_img[i][j];
      A_[0][idx] = p.r;
      A_[1][idx] = p.g;
      A_[2][idx] = p.b;
      ++idx;
    }
  }
  
  haar_.Transform(A_[0], A_[1], A_[2]);

  SigPtr sig = SigPtr(new Signature);
  sig->id = s_guid.NewId();
  sig->width = image.width();
  sig->height = image.height();
  
  for (int c = 0; c < 3; ++c) {
    haar_.GetCoefs(A_[c], sig->sigs[c], &sig->avgs[c]);
  }

  return sig;
}

} // oe
